package com.shujia.mllib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo05ImageTrain {
  def main(args: Array[String]): Unit = {

    val spark: SparkSession = SparkSession
      .builder()
      .appName("Demo05ImageTrain")
      .master("local[*]")
      .config("spark.sql.shuffle.partitions", "8")
      .getOrCreate()

    import spark.implicits._
    import org.apache.spark.sql.functions._

    // 1、加载特征工程处理好的数据
    val imageDataDF: DataFrame = spark
      .read
      .format("libsvm")
      .option("numFeatures", 28 * 28)
      .load("Spark/data/mllib/imageFeatures")

    // 2、切分数据集
    val arr: Array[Dataset[Row]] = imageDataDF.randomSplit(Array(0.8, 0.2))
    val trainDF: Dataset[Row] = arr(0)
    val testDF: Dataset[Row] = arr(1)

    // 3、选择合适的模型 ==> 逻辑回归
    val logisticRegression: LogisticRegression = new LogisticRegression()
      .setMaxIter(10) // 最大的迭代次数
      .setFitIntercept(true) // 设置是否有截距

    // 4、使用训练集进行训练
    val logisticRegressionModel: LogisticRegressionModel = logisticRegression.fit(trainDF)

    // 5、使用测试集进行评估
    val testTranDF: DataFrame = logisticRegressionModel.transform(testDF)

    testTranDF.show(100, truncate = false)

    // 6、计算模型的准确率
    testTranDF
      .withColumn("isEqual", when($"label" === $"prediction", 1).otherwise(0))
      .select(sum($"isEqual") / count("*") as "准确率")
      .show()

    // 7、保存模型
    logisticRegressionModel.save("Spark/data/mllib/image")
  }

}
